#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import logging
import requests
from typing import Dict, List, Any, Optional

from vector_store.vector_store_interface import VectorStoreInterface
from vector_store.sqlite_vector_store import SQLiteVectorStore

# 获取logger
logger = logging.getLogger(__name__)

class NodeVectorProcessor:
    """
    节点向量处理器，用于处理知识图谱中class和function节点的向量化
    """

    def __init__(self, vector_store: Optional[VectorStoreInterface] = None):
        """
        初始化节点向量处理器

        Args:
            vector_store: 向量存储对象，如果为None则创建SQLiteVectorStore实例
        """
        self.vector_store = vector_store
        
        # 获取嵌入模型配置

        # 如果向量存储没有api_key属性，从全局配置获取
        from config.config import config
        embedding_config = config.get_embedding_config()
        self.api_key = embedding_config.get('api_key')
        self.model_name = embedding_config.get('model_name')
        self.api_base = embedding_config.get('api_base')
        
        logger.info("节点向量处理器初始化成功")
        logger.debug(f"使用的嵌入模型: {self.model_name}")
        logger.debug(f"使用的API基础URL: {self.api_base}")

    def get_embedding(self, text: str) -> List[float]:
        """
        获取文本的嵌入向量

        Args:
            text: 需要嵌入的文本

        Returns:
            嵌入向量
        """
        try:
            headers = {
                "Content-Type": "application/json",
                "Authorization": f"Bearer {self.api_key}"
            }

            payload = {
                "model": self.model_name,
                "input": text,
                "encoding_format": "float"
            }

            response = requests.post(
                f"{self.api_base}/embeddings",
                headers=headers,
                json=payload
            )

            if response.status_code == 200:
                result = response.json()
                return result["data"][0]["embedding"]
            else:
                logger.error(f"获取嵌入向量失败: {response.status_code} {response.text}")
                # 返回一个空向量作为回退方案
                return [0.0] * 1024  # bge-m3的向量维度
        except Exception as e:
            logger.error(f"调用嵌入API失败: {str(e)}")
            return [0.0] * 1024  # 返回空向量

    def process_class_node(self, node_info: Dict[str, Any]) -> bool:
        """
        处理类节点，提取文本并保存向量

        Args:
            node_info: 类节点信息

        Returns:
            处理是否成功
        """
        try:
            if "full_name" not in node_info:
                logger.warning("类节点缺少full_name属性，无法处理")
                return False

            if "node" not in node_info:
                logger.warning("类节点缺少node属性，无法处理")
                return False

            # 从AST节点直接获取文本内容
            # try:
            #     node_text = node_info["node"].text.decode("utf-8")
            # except Exception as e:
            #     logger.warning(f"无法解码节点文本: {str(e)}")
                # 回退到构建简单文本表示
            node_text = f"{node_info.get('type', 'class')} {node_info.get('full_name', '')}"

            # 提取元数据
            metadata = {
                "name": node_info.get("name", ""),
                "full_name": node_info.get("full_name", ""),
                "visibility": node_info.get("visibility", "default"),
                "class_type": node_info.get("type", "class"),
                "line_start": node_info.get("line_start", 0),
                "line_end": node_info.get("line_end", 0)
            }

            # 添加继承信息
            if "extends" in node_info:
                metadata["extends"] = node_info["extends"]

            # 添加实现接口信息
            if "implements" in node_info and isinstance(node_info["implements"], list):
                metadata["implements"] = node_info["implements"]

            # 生成向量
            vector = self.get_embedding(node_text)

            # 存储向量
            return self.vector_store.store_node_vector(
                vid=node_info["full_name"],
                node_type="class",
                node_text=node_text,
                vector=vector,
                metadata=metadata
            )
        except Exception as e:
            logger.error(f"处理类节点失败: {str(e)}")
            return False

    def process_function_node(self, node_info: Dict[str, Any]) -> bool:
        """
        处理函数节点，提取文本并保存向量

        Args:
            node_info: 函数节点信息

        Returns:
            处理是否成功
        """
        try:
            if "full_name" not in node_info:
                logger.warning("函数节点缺少full_name属性，无法处理")
                return False

            if "node" not in node_info:
                logger.warning("函数节点缺少node属性，无法处理")
                return False

            # 从AST节点直接获取文本内容
            # try:
            #     node_text = node_info["node"].text.decode("utf-8")
            # except Exception as e:
            #     logger.warning(f"无法解码节点文本: {str(e)}")
            #     # 回退到构建简单文本表示
            node_text = (f"{node_info.get('return_type', 'void')} "
                             f"{node_info.get('signature', node_info.get('name', ''))}")

            # 提取元数据
            metadata = {
                "name": node_info.get("name", ""),
                "full_name": node_info.get("full_name", ""),
                "visibility": node_info.get("visibility", "default"),
                "is_static": node_info.get("is_static", False),
                "is_constructor": node_info.get("is_constructor", False),
                "return_type": node_info.get("return_type", "void"),
                "line_start": node_info.get("line_start", 0),
                "line_end": node_info.get("line_end", 0)
            }

            # 添加参数信息
            if "parameters" in node_info and node_info["parameters"]:
                metadata["parameters"] = node_info["parameters"]

            # 添加签名信息
            if "signature" in node_info:
                metadata["signature"] = node_info["signature"]

            # 生成向量
            vector = self.get_embedding(node_text)

            # 存储向量
            return self.vector_store.store_node_vector(
                vid=node_info["full_name"],
                node_type="function",
                node_text=node_text,
                vector=vector,
                metadata=metadata
            )
        except Exception as e:
            logger.error(f"处理函数节点失败: {str(e)}")
            return False

    def process_annotation_node(self, node_info: Dict[str, Any]) -> bool:
        """
        处理注解节点，提取文本并保存向量

        Args:
            node_info: 注解节点信息

        Returns:
            处理是否成功
        """
        try:
            if "full_name" not in node_info:
                logger.warning("注解节点缺少full_name属性，无法处理")
                return False

            # 尝试从AST节点获取文本
            node_text = ""
            if "node" in node_info:
                try:
                    node_text = node_info["node"].text.decode("utf-8")
                except Exception as e:
                    logger.warning(f"无法从AST节点获取文本: {str(e)}")

            # 如果无法从AST节点获取文本，则构建基本文本表示
            if not node_text:
                node_text = f"Annotation {node_info.get('name', '')} {node_info.get('full_name', '')}"

            # 如果有内容，添加到文本中
            if "content" in node_info and node_info["content"]:
                try:
                    # 尝试解压内容
                    from parser.utils import ContentCompressor
                    decompressed_content = ContentCompressor.decompress(node_info["content"])
                    node_text += f"\nContent: {decompressed_content}"
                except Exception as e:
                    logger.warning(f"无法解压注解内容: {str(e)}")
                    node_text += f"\nContent: {node_info['content']}"

            # 提取元数据
            metadata = {
                "name": node_info.get("name", ""),
                "full_name": node_info.get("full_name", ""),
                "type": node_info.get("type", "ANNOTATION"),
                "line_start": node_info.get("line_start", 0),
                "line_end": node_info.get("line_end", 0)
            }

            # 生成向量
            vector = self.get_embedding(node_text)

            # 存储向量
            return self.vector_store.store_node_vector(
                vid=node_info["full_name"],
                node_type="annotation",
                node_text=node_text,
                vector=vector,
                metadata=metadata
            )
        except Exception as e:
            logger.error(f"处理注解节点失败: {str(e)}")
            return False

    def process_entities(self, entities: Dict[str, List[Dict[str, Any]]]) -> Dict[str, int]:
        """
        处理实体集合，提取类、函数和注解节点并保存向量

        Args:
            entities: 按类型分组的实体列表

        Returns:
            处理统计信息
        """
        stats = {
            "total": 0,
            "class": 0,
            "function": 0,
            "annotation": 0,
            "succeeded": 0,
            "failed": 0
        }

        # 处理类节点
        for class_node in entities.get("class", []):
            stats["total"] += 1
            stats["class"] += 1

            success = self.process_class_node(class_node)
            if success:
                stats["succeeded"] += 1
            else:
                stats["failed"] += 1

        # 处理函数节点
        for function_node in entities.get("function", []):
            stats["total"] += 1
            stats["function"] += 1

            success = self.process_function_node(function_node)
            if success:
                stats["succeeded"] += 1
            else:
                stats["failed"] += 1

        # 处理注解节点
        for annotation_node in entities.get("annotations", []):
            # 只处理type=ANNOTATION的注解节点
            if annotation_node.get("type") == "ANNOTATION":
                stats["total"] += 1
                stats["annotation"] += 1

                success = self.process_annotation_node(annotation_node)
                if success:
                    stats["succeeded"] += 1
                else:
                    stats["failed"] += 1

        return stats

    def search_similar_nodes(self, query_text: str, node_type: Optional[str] = None, top_k: int = 10) -> List[Dict[str, Any]]:
        """
        搜索与查询文本最相似的节点

        Args:
            query_text: 查询文本
            node_type: 节点类型过滤 ('class' 或 'function'，None表示不过滤)
            top_k: 返回结果数量

        Returns:
            相似节点列表，按相似度降序排序
        """
        # 生成查询向量
        query_vector = self.get_embedding(query_text)

        # 使用向量存储执行相似搜索
        return self.vector_store.search_similar_vectors(query_vector, node_type, top_k)

    def close(self) -> None:
        """关闭资源"""
        if self.vector_store:
            self.vector_store.close()